package ai.neuralnet;

import java.util.ArrayList;

public class ClassificationDataSet
{
    private ArrayList<double[]> inputValues;

    private ArrayList<double[]> outputValues;

    public ClassificationDataSet()
    {
        inputValues = new ArrayList<double[]>();
        outputValues = new ArrayList<double[]>();
    }


    public int getCategory(int index)
    {
        for (int i = 0; i < outputValues.get(index).length; i++)
        {
            if (outputValues.get(index)[i] == 1) return i;
        }

        return -1;
    }

    public int getPointDimension()
    {
        return inputValues.get(0).length;
    }

    public int getCategoryCount()
    {
        return outputValues.get(0).length;
    }

    public void addAll(ClassificationDataSet set)
    {
        if (set.getPointDimension() != getPointDimension())
        {
            System.err.println("Error: cannot add points with different dimensions to ClassificationDataSet");
            return;
        }

        for (int i = 0; i < set.size(); i++)
        {
            addInputOutputPair(set.getInputValues().get(i), set.getOutputValues().get(i));
        }
    }

    public int[] getPredictions(NeuralNetwork network)
    {
        int[] predictions = new int[inputValues.size()];
        for (int i = 0; i < inputValues.size(); i++)
        {
            predictions[i] = predictedCategory(network, inputValues.get(i))[0];
        }

        return predictions;
    }

    public void addInputOutputPair(double[] input, double[] output)
    {
        inputValues.add(input);
        outputValues.add(output);
    }

    public int[] predictedCategory(NeuralNetwork network, double[] input)
    {
        network.propagateInputs(input);
        double[] output = network.getOutputs();
        int index = -1;
        double max = Double.NEGATIVE_INFINITY, nextMax = Double.NEGATIVE_INFINITY;
        for (int i = 0; i < output.length; i++)
        {
            if (output[i] >= max)
            {
                nextMax = max;
                max = output[i];
                index = i;
            }

            if (output[i] < max && output[i] > nextMax)
            {
                nextMax = output[i];
            }
        }

        return new int[]{index, max > 0.75 && nextMax < 0.25 ? index : -1};
        // return ;
    }

    public ArrayList<double[]> getInputValues()
    {
        return inputValues;
    }

    public ArrayList<double[]> getOutputValues()
    {
        return outputValues;
    }

    public double[] getAccuracy(NeuralNetwork network)
    {
        int bgCorrect = 0, cCorrect = 0;
        for (int i = 0; i < inputValues.size(); i++)
        {
            int cat[] = predictedCategory(network, inputValues.get(i));

            if (outputValues.get(i)[cat[0]] == 1) bgCorrect++;
            if (cat[1] != -1 && outputValues.get(i)[cat[1]] == 1) cCorrect++;
        }

        return new double[]{bgCorrect / (double) outputValues.size(), cCorrect / (double) outputValues.size()};
    }

    public int size()
    {
        return inputValues.size();
    }


    public double getAccuracy(int[] predictions)
    {
        int correct = 0;
        for (int i = 0; i < inputValues.size(); i++)
        {
            if (outputValues.get(i)[predictions[i]] == 1) correct++;
        }

        return correct / (double) outputValues.size();
    }
}
